# -*- coding: utf-8 -*-
"""
Created on Thu Oct 10 13:59:03 2013

@author: jkwong
"""

#PBAR_LiningUpRadZspecImages.py


import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage
import copy
import glob

# list of zpsec detector names - Z1 to Z136
ZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137)])
# list of bad detector numbers starting at 1
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
# list of good detector numbers starting at 1
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])
# list of bad detector names
badZspecNameList = ZspecNameList[(badZspecList-1)]
# list of good detector names
goodZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecIndices = badZspecList - 1
goodZspecIndices = goodZspecList - 1

badZspecMask = np.zeros(len(ZspecNameList))
badZspecMask[badZspecIndices] = True

goodZspecMask = np.zeros(len(ZspecNameList))
goodZspecMask[goodZspecIndices] = True


#################################
## DEFINE FILE NAMES
#
## Configuration 20
#
#cargoConfig = '20A'
#
## Zspec Files
#fullFilenameList = []
#fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\4761-FDFC-All2148.npy')
#fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\4763-FDFC-All2151.npy')
#
#filenameList = []
#for f in fullFilenameList:
#    a, b = os.path.split(f)
#    filenameList.append(b)
#
## Corresponding cargo images
#filenameCargoList = []
#filenameCargoList.append('PBAR-20130902007.cargoimage')
#filenameCargoList.append('PBAR-20130902009.cargoimage')
#
#fullFilenameCargoList = []
#fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargoimage')
#fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargoimage')
#
## Marker files
#markerfullFilenameCargoList = []
#markerfullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargomarker')
#markerfullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargomarker')
#
# some number related to the datasets
acquisitionTime = np.array([1/60., 1/60.])
airSection = np.array([0, 175])

numberTimeSlices = 2000
preCargoTimeSlices = 70

cargoConfigList = ['7A', '7B', '9A', '9B']


acquisitionTimeList = []
zspecScanNumberList = []
cargoScanNumberList = []

fullFilenameList = []
filenameList = []

filenameCargoList = []
fullFilenameCargoList = []
markerfullFilenameCargoList = []


for (index, cargoConfig) in enumerate(cargoConfigList):
    
    datapath = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScans\%s\%s' %(cargoConfig[0:(len(cargoConfig)-1)], cargoConfig[-1])
    
    # find the zspec scan number
    temp = glob.glob(os.path.join(datapath, '*.npy'))
    a, filename = os.path.split(temp[0])
    zspecScanNumberList.append(filename[0:4])
    filenameList.append(filename)
    
    # get the cargo image name
    temp = glob.glob(os.path.join(datapath, '*.cargoimage'))
    a, filename = os.path.split(temp[0])
    cargoScanNumberList.append(filename[0:16])
    filenameCargoList.append(filename)

    fullFilenameList.append(os.path.join(datapath, filenameList[index]))
    fullFilenameCargoList.append(os.path.join(datapath, filenameCargoList[index]))
    markerfullFilenameCargoList.append(os.path.join(datapath, filenameCargoList[index].replace('cargoimage', 'cargomarker')))
    acquisitionTimeList.append(1/60.)


#fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\4761-FDFC-All2148.npy')
#fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\4763-FDFC-All2151.npy')
#
#
## Corresponding cargo images
#filenameCargoList.append('PBAR-20130902007.cargoimage')
#filenameCargoList.append('PBAR-20130902009.cargoimage')
#
#fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargoimage')
#fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargoimage')
## Marker files
#markerfullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargomarker')
#markerfullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargomarker')
#
#
#
#for f in fullFilenameList:
#    a, b = os.path.split(f)
#    filenameList.append(b)
#
#for f in fullFilenameCargoList:
#    a, b = os.path.split(f)
#    filenameCargoList.append(b)

######################
## READ IN DATA

# adjust images as read in

# Read in marker files
markerList = []
for (index, f) in enumerate(markerfullFilenameCargoList):
    markerList.append(PBAR_Cargo.ReadCargoMarker(f))

markerStandardWidthList = []

energyList = []
dat = []
datStandardWidth = []
startBinZspecList = []
for (index, fullFilename) in enumerate(fullFilenameList):
    (a, b) = PBAR_Zspec.ReadZspecBasicScanNumpy(fullFilename)
    energyList.append(a)
    dat.append(b)
    # reverse, limit to 200 bins, reverse back
#    datStandardWidth.append(b[-1::-1,:,:][0:2000,:,:][-1::-1,:,:])
    
#    datStandardWidth.append(b[-1::-1,:,:][0:2000,:,:][-1::-1,:,:])
    startBinZspec = PBAR_Zspec.FindBasicZspecStart(b)
    startBinZspecList.append(startBinZspec)
    
    temp = b[(startBinZspec-preCargoTimeSlices):,:][0:numberTimeSlices,:,:] # this could have length less than numberTimeSlices    
 
   # This ensure that we have standard size images    
    temp2 = np.zeros((numberTimeSlices, b.shape[1], b.shape[2]))
    if temp.shape[0] < numberTimeSlices:
        temp2[0:temp.shape[0],:,:] = temp
    else:
        temp2 = temp
    datStandardWidth.append(temp2)

del a, b

# read in cargo images
numberZspecDetectors = dat[0].shape[1]
datCargo = []
datCargoStandardWidth = []
startBinCargoList = []
#datCargoStandardWidth2 = []

for (index, f) in enumerate(fullFilenameCargoList):
    (A,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(f)
    datCargo.append(A)
    # flip, cut to 2000, flip again
    # 1) throw away last 4 entries, 2) reverse, 3) limit length to 2000 slices
    # 4) Reverse back to original time order
#    datCargoStandardWidth.append(A[-5::-1,:][0:2000,:][-1::-1,:])
    
#    temp = np.zeros((A.shape[0], numberZspecDetectors))
#    for zspecChannel in xrange(numberZspecDetectors):
#        cargoChannel1 = 1 + zspecChannel*4
#        cargoChannel2 = 2 + zspecChannel*4        
#        temp[:,zspecChannel] = (A[:,cargoChannel1] + A[:,cargoChannel2])/2  # check if overflow issues
#    datCargoStandardWidth.append(temp)
    # this should do the same thing
    indices1 = 1 + 4 * np.arange(136) 
    indices2 = 2 + 4* np.arange(136)
    temp = (A[:,indices1] + A[:,indices2])/2
    # line up in time
#    datCargoStandardWidth.append(temp[-5::-1,:][0:2000][-1::-1,:])
    
    startBinCargo = PBAR_Cargo.FindCargoStart(A)
    startBinCargoList.append(startBinCargo)
    
    # cut section in front and limit to numberTimeSlices number of time slices
    temp = temp[(startBinCargo-preCargoTimeSlices):,:][0:numberTimeSlices,:]
    
    # This ensure that we have standard size images
    # excess is filled with zeros
    temp2 = np.zeros((numberTimeSlices, numberZspecDetectors))
    if temp.shape[0] < numberTimeSlices:
        temp2[0:temp.shape[0],:] = temp
    else:
        temp2 = temp[0:numberTimeSlices,:]
        
    datCargoStandardWidth.append(temp2)    
    
del A,bpp,formatt,flag,low1,high1,low2,high2
del indices1, indices2, temp, temp2

# Modify the marker files

# {'color': u'0xFFFF8800',
#  'done': u'0',
#  u'left': {'x': 387.0, 'y': 231.0},
#  'locked': u'0',
#  'rec_bottom': 242.0,
#  'rec_left': 403.0,
#  'rec_right': 427.0,
#  'rec_top': 223.0,
#  u'right': {'x': 441.0, 'y': 231.0},
#  'shape': u'rectangle',
#  'suspicious': u'0',
#  'target': u'F4',
#  'x': 414.0,
#  'y': 231.0}]

markerStandardWidthList = []
for (index, markers) in enumerate(markerList):
    # first copy over
    temp = copy.deepcopy(markers)
    
    # adjust the x values
    offset = (startBinCargo - preCargoTimeSlices)
    for i in xrange(len(temp)):
        temp[i]['rec_left'] = temp[i]['rec_left'] - offset
        temp[i]['rec_right'] = temp[i]['rec_right'] - offset
        temp[i]['x'] = temp[i]['x'] - offset
        # see if 'left and 'right exist
        if 'left' in temp[i]:
            temp[i]['left']['x'] = temp[i]['left']['x'] - offset
        if 'right' in temp[i]:
            temp[i]['right']['x'] = temp[i]['right']['x'] - offset
    # adjust the y values
    for i in xrange(len(temp)):
        temp[i]['rec_top'] = np.round((temp[i]['rec_top'] - 1.5) / 4.0)
        temp[i]['rec_bottom'] = np.round((temp[i]['rec_bottom'] - 1.5) / 4.0)
        temp[i]['y'] = np.round((temp[i]['y'] - 1.5) / 4.0)
        # see if 'left and 'right exist
        if 'left' in temp[i]:
            temp[i]['left']['y'] = np.round((temp[i]['left']['y'] - 1.5) / 4.0)
        if 'right' in temp[i]:
            temp[i]['right']['y'] = np.round((temp[i]['right']['y'] - 1.5) / 4.0)
    markerStandardWidthList.append(copy.copy(temp))


###################
## PLOTS 

## plot raw unadjusted images, both
# for zspec plot the total counts
index = 1
logIntensity = True
plotMarkers = True

plt.figure()
plt.grid()

if logIntensity:
    intensity = np.log(datCargo[index].T)
else:
    intensity = datCargo[index].T

plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2400))

if plotMarkers:
    marker = markerList[index]
    for i, mark in enumerate( marker):
        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
        plt.plot(x, y, 'r')
        plt.plot(mark['x'], mark['y'], 'xr')
        if 'left' in mark:
            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
        if 'right' in mark:
            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
plt.figure()
plt.grid()

if logIntensity:
    intensity = np.log(dat[index].sum(axis = 2).T)
else:
    intensity = dat[index].sum(axis = 2).T

plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameList[index])
plt.xlim((0, 2400))


## plot images that have been set to standard size
# for zspec plot the total counts
index = 0
logIntensity = 1
plotMarkers = True
plt.figure()
plt.grid()

if logIntensity:
    intensity = np.log(datCargoStandardWidth[index].T)
else:
    intensity = datCargoStandardWidth[index].T
    
plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2400))

if plotMarkers:
    marker = markerStandardWidthList[index]
    for i, mark in enumerate( marker):
        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
        plt.plot(x, y, 'r')
        plt.plot(mark['x'], mark['y'], 'xr')
        if 'left' in mark:
            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
        if 'right' in mark:
            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
    
plt.figure()
plt.grid()

plt.imshow((datStandardWidth[index].sum(axis = 2).T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameList[index])
plt.xlim((0, 2400))

if plotMarkers:
    marker = markerStandardWidthList[index]
    for i, mark in enumerate( marker):
        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
        plt.plot(x, y, 'r')
        plt.plot(mark['x'], mark['y'], 'xr')
        if 'left' in mark:
            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
        if 'right' in mark:
            plt.plot(mark['right']['x'], mark['right']['y'], 'go')


# PLOT BOTH IMAGES WITH BEST OVERLAP POSSIBLE
threshold = 1.5e9

plt.figure()
plt.grid()

index = 1
plt.imshow(log(datCargoStandardWidth[index].T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))

# zspec image
plt.figure()
plt.grid()

discrimName = 'binSTD_binMean'
discrimName = 'binSTD'
discrimName = 'count'
#discrimName = 'multibin_20_ratio'
index = 1
plt.imshow(discrim[index][discrimName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmin = 0, vmax = 500)
#plt.imshow(discrim[index][discrimName][-1::-1,1:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(discrimName)
plt.xlim((0, 2200))